# rag/store.py
# from datetime import time
import time
import faiss_rag
import pickle
import numpy as np
import os
import faiss
from typing import List, Dict
from faiss_rag.embed import get_embedding
from faiss_rag.config import FAISS_INDEX_PATH, VECTOR_DIM


class EntertainmentRAGStore:
    def __init__(self):
        self.index_path = FAISS_INDEX_PATH
        self.index_file = os.path.join(self.index_path, "faiss_rag.index")
        self.meta_file = os.path.join(self.index_path, "metadata.pkl")

        # 创建目录
        os.makedirs(self.index_path, exist_ok=True)

        if os.path.exists(self.index_file):
            self.index = faiss.read_index(self.index_file)
        else:
            self.index = faiss.IndexFlatL2(VECTOR_DIM)  # 空索引

        if os.path.exists(self.meta_file):
            with open(self.meta_file, 'rb') as f:
                self.metadata = pickle.load(f)
        else:
            self.metadata = []

    def add_texts(self, texts: List[str]):
        """批量添加娱乐/情感类文本到 FAISS"""
        vectors = []
        new_meta = []

        for text in texts:
            vec = get_embedding(text)
            vectors.append(vec)
            new_meta.append({"text": text, "added_at": time.time()})

        # 转为 numpy array
        vectors_np = np.array(vectors, dtype=np.float32)

        # 添加到索引
        self.index.add(vectors_np)
        self.metadata.extend(new_meta)

        # 持久化
        faiss.write_index(self.index, self.index_file)
        with open(self.meta_file, 'wb') as f:
            pickle.dump(self.metadata, f)

        print(f"✅ 已添加 {len(texts)} 条娱乐/情感新闻到 FAISS")

    def __len__(self):
        return self.index.ntotal

if __name__ == '__main__':
    store = EntertainmentRAGStore()
    store.add_texts(['aaaaabbbbb'])